import torch
import torch.nn as nn
import torch.nn.functional as F
from einops.layers.torch import Rearrange
from torch import Tensor
class CNNFeatureExtractor(nn.Module):
    def __init__(self, emb_size=20):
        super().__init__()
        self.shallownet = nn.Sequential(
            nn.Conv2d(1, 20, (1, 25), (1, 1)),
            nn.Conv2d(20, 20, (22, 1), (1, 1)),
            nn.BatchNorm2d(20),
            nn.ELU(),
            nn.AvgPool2d((1, 45), (1, 15)),
            nn.Dropout(0.5),
            nn.Conv2d(20, emb_size, (1, 1), stride=(1, 1)),
            Rearrange('b e (h) (w) -> b (h w) e')
        )

    def forward(self, x):
        return self.shallownet(x)

class LSTMBlock(nn.Module):
    def __init__(self, input_dim, hidden_size, num_layers, dropout):
        super().__init__()
        self.lstm = nn.LSTM(input_dim, hidden_size, num_layers, batch_first=True, dropout=dropout)

    def forward(self, x):
        return self.lstm(x)[0][:, -1, :]

class CNNLSTM(nn.Module):
    def __init__(self, emb_size=20, hidden_size=64, num_layers=2, dropout=0.7, n_classes=2,device=None):
        super().__init__()
        self.feature_extractor = CNNFeatureExtractor(emb_size=emb_size)
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.lstm_block = LSTMBlock(emb_size, hidden_size, num_layers, dropout)
        self.classifier = nn.Linear(hidden_size, n_classes)

    def forward(self, x, return_features=False):
        features = self.feature_extractor(x)
        lstm_output = self.lstm_block(features)
        if return_features:
            return self.classifier(lstm_output)
        return self.classifier(lstm_output)
